/********************************************************************************/
/*										*/
/*			     The RSA key cache 					*/
/*			     Written by Ken Goldman				*/
/*		       IBM Thomas J. Watson Research Center			*/
/*										*/
/*  Licenses and Notices							*/
/*										*/
/*  1. Copyright Licenses:							*/
/*										*/
/*  - Trusted Computing Group (TCG) grants to the user of the source code in	*/
/*    this specification (the "Source Code") a worldwide, irrevocable, 		*/
/*    nonexclusive, royalty free, copyright license to reproduce, create 	*/
/*    derivative works, distribute, display and perform the Source Code and	*/
/*    derivative works thereof, and to grant others the rights granted herein.	*/
/*										*/
/*  - The TCG grants to the user of the other parts of the specification 	*/
/*    (other than the Source Code) the rights to reproduce, distribute, 	*/
/*    display, and perform the specification solely for the purpose of 		*/
/*    developing products based on such documents.				*/
/*										*/
/*  2. Source Code Distribution Conditions:					*/
/*										*/
/*  - Redistributions of Source Code must retain the above copyright licenses, 	*/
/*    this list of conditions and the following disclaimers.			*/
/*										*/
/*  - Redistributions in binary form must reproduce the above copyright 	*/
/*    licenses, this list of conditions	and the following disclaimers in the 	*/
/*    documentation and/or other materials provided with the distribution.	*/
/*										*/
/*  3. Disclaimers:								*/
/*										*/
/*  - THE COPYRIGHT LICENSES SET FORTH ABOVE DO NOT REPRESENT ANY FORM OF	*/
/*  LICENSE OR WAIVER, EXPRESS OR IMPLIED, BY ESTOPPEL OR OTHERWISE, WITH	*/
/*  RESPECT TO PATENT RIGHTS HELD BY TCG MEMBERS (OR OTHER THIRD PARTIES)	*/
/*  THAT MAY BE NECESSARY TO IMPLEMENT THIS SPECIFICATION OR OTHERWISE.		*/
/*  Contact TCG Administration (admin@trustedcomputinggroup.org) for 		*/
/*  information on specification licensing rights available through TCG 	*/
/*  membership agreements.							*/
/*										*/
/*  - THIS SPECIFICATION IS PROVIDED "AS IS" WITH NO EXPRESS OR IMPLIED 	*/
/*    WARRANTIES WHATSOEVER, INCLUDING ANY WARRANTY OF MERCHANTABILITY OR 	*/
/*    FITNESS FOR A PARTICULAR PURPOSE, ACCURACY, COMPLETENESS, OR 		*/
/*    NONINFRINGEMENT OF INTELLECTUAL PROPERTY RIGHTS, OR ANY WARRANTY 		*/
/*    OTHERWISE ARISING OUT OF ANY PROPOSAL, SPECIFICATION OR SAMPLE.		*/
/*										*/
/*  - Without limitation, TCG and its members and licensors disclaim all 	*/
/*    liability, including liability for infringement of any proprietary 	*/
/*    rights, relating to use of information in this specification and to the	*/
/*    implementation of this specification, and TCG disclaims all liability for	*/
/*    cost of procurement of substitute goods or services, lost profits, loss 	*/
/*    of use, loss of data or any incidental, consequential, direct, indirect, 	*/
/*    or special damages, whether under contract, tort, warranty or otherwise, 	*/
/*    arising in any way out of use or reliance upon this specification or any 	*/
/*    information herein.							*/
/*										*/
/*  (c) Copyright IBM Corp. and others, 2016 - 2023				*/
/*										*/
/********************************************************************************/

//** Introduction
// This file contains the functions to implement the RSA key cache that can be used
// to speed up simulation.
//
// Only one key is created for each supported key size and it is returned whenever
// a key of that size is requested.
//
// If desired, the key cache can be populated from a file. This allows multiple
// TPM to run with the same RSA keys. Also, when doing simulation, the DRBG will
// use preset sequences so it is not too hard to repeat sequences for debug or
// profile or stress.
//
// When the key cache is enabled, a call to CryptRsaGenerateKey() will call the
// GetCachedRsaKey(). If the cache is enabled and populated, then the cached key
// of the requested size is returned. If a key of the requested size is not
// available, the no key is loaded and the requested key will need to be generated.
// If the cache is not populated, the TPM will open a file that has the appropriate
// name for the type of keys required (CRT or no-CRT). If the file is the right
// size, it is used. If the file doesn't exist or the file does not have the correct
// size, the TMP will populate the cache with new keys of the required size and
// write the cache data to the file so that they will be available the next time.
//
// Currently, if two simulations are being run with TPM's that have different RSA
// key sizes (e.g,, one with 1024 and 2048 and another with 2048 and 3072, then the
// files will not match for the both of them and they will both try to overwrite
// the other's cache file. I may try to do something about this if necessary.

//** Includes, Types, Locals, and Defines

#include "Tpm.h"

#if USE_RSA_KEY_CACHE

#  include <stdio.h>
#  include "RsaKeyCache_fp.h"

#  if CRT_FORMAT_RSA == YES
#    define CACHE_FILE_NAME "RsaKeyCacheCrt.data"
#  else
#    define CACHE_FILE_NAME "RsaKeyCacheNoCrt.data"
#  endif

typedef struct _RSA_KEY_CACHE_
{
    TPM2B_PUBLIC_KEY_RSA  publicModulus;
    TPM2B_PRIVATE_KEY_RSA privateExponent;
} RSA_KEY_CACHE;

// Determine the number of RSA key sizes for the cache
TPMI_RSA_KEY_BITS SupportedRsaKeySizes[] = {
#  if RSA_1024
    1024,
#  endif
#  if RSA_2048
    2048,
#  endif
#  if RSA_3072
    3072,
#  endif
#  if RSA_4096
    4096,
#  endif
    0};

#  define RSA_KEY_CACHE_ENTRIES (RSA_1024 + RSA_2048 + RSA_3072 + RSA_4096)

// The key cache holds one entry for each of the supported key sizes
RSA_KEY_CACHE s_rsaKeyCache[RSA_KEY_CACHE_ENTRIES];
// Indicates if the key cache is loaded. It can be loaded and enabled or disabled.
BOOL s_keyCacheLoaded = 0;

// Indicates if the key cache is enabled
int s_rsaKeyCacheEnabled = FALSE;

//*** RsaKeyCacheControl()
// Used to enable and disable the RSA key cache.
LIB_EXPORT void RsaKeyCacheControl(int state)
{
    s_rsaKeyCacheEnabled = state;
}

//*** InitializeKeyCache()
// This will initialize the key cache and attempt to write it to a file for later
// use.
//  Return Type: BOOL
//      TRUE(1)         success
//      FALSE(0)        failure
static BOOL InitializeKeyCache(TPMT_PUBLIC*    publicArea,
			       TPMT_SENSITIVE* sensitive,
			       RAND_STATE* rand  // IN: if not NULL, the deterministic
			       //     RNG state
			       )
{
    int          index;
    TPM_KEY_BITS keySave = publicArea->parameters.rsaDetail.keyBits;
    BOOL         OK      = TRUE;
    //
    s_rsaKeyCacheEnabled = FALSE;
    for(index = 0; OK && index < RSA_KEY_CACHE_ENTRIES; index++)
	{
	    publicArea->parameters.rsaDetail.keyBits = SupportedRsaKeySizes[index];
	    OK = (CryptRsaGenerateKey(publicArea, sensitive, rand) == TPM_RC_SUCCESS);
	    if(OK)
		{
		    s_rsaKeyCache[index].publicModulus   = publicArea->unique.rsa;
		    s_rsaKeyCache[index].privateExponent = sensitive->sensitive.rsa;
		}
	}
    publicArea->parameters.rsaDetail.keyBits = keySave;
    s_keyCacheLoaded                         = OK;
#  if SIMULATION && USE_RSA_KEY_CACHE && USE_KEY_CACHE_FILE
    if(OK)
	{
	    FILE*       cacheFile;
	    const char* fn = CACHE_FILE_NAME;

#    if defined     _MSC_VER
	    if(fopen_s(&cacheFile, fn, "w+b") != 0)
#    else
		cacheFile = fopen(fn, "w+b");
	    if(NULL == cacheFile)
#    endif
		{
		    printf("Can't open %s for write.\n", fn);
		}
	    else
		{
		    fseek(cacheFile, 0, SEEK_SET);
		    if(fwrite(s_rsaKeyCache, 1, sizeof(s_rsaKeyCache), cacheFile)
		       != sizeof(s_rsaKeyCache))
			{
			    printf("Error writing cache to %s.", fn);
			}
		}
	    if(cacheFile)
		fclose(cacheFile);
	}
#  endif
    return s_keyCacheLoaded;
}

//*** KeyCacheLoaded()
// Checks that key cache is loaded.
//  Return Type: BOOL
//      TRUE(1)         cache loaded
//      FALSE(0)        cache not loaded
static BOOL KeyCacheLoaded(TPMT_PUBLIC*    publicArea,
			   TPMT_SENSITIVE* sensitive,
			   RAND_STATE*     rand  // IN: if not NULL, the deterministic
			   //     RNG state
			   )
{
#  if SIMULATION && USE_RSA_KEY_CACHE && USE_KEY_CACHE_FILE
    if(!s_keyCacheLoaded)
	{
	    FILE*       cacheFile;
	    const char* fn = CACHE_FILE_NAME;
#    if defined     _MSC_VER && 1
	    if(fopen_s(&cacheFile, fn, "r+b") == 0)
#    else
		cacheFile = fopen(fn, "r+b");
	    if(NULL != cacheFile)
#    endif
		{
		    fseek(cacheFile, 0L, SEEK_END);
		    if(ftell(cacheFile) == sizeof(s_rsaKeyCache))
			{
			    fseek(cacheFile, 0L, SEEK_SET);
			    s_keyCacheLoaded =
				(fread(&s_rsaKeyCache, 1, sizeof(s_rsaKeyCache), cacheFile)
				 == sizeof(s_rsaKeyCache));
			}
		    fclose(cacheFile);
		}
	}
#  endif
    if(!s_keyCacheLoaded)
	s_rsaKeyCacheEnabled = InitializeKeyCache(publicArea, sensitive, rand);
    return s_keyCacheLoaded;
}

//*** GetCachedRsaKey()
//  Return Type: BOOL
//      TRUE(1)         key loaded
//      FALSE(0)        key not loaded
BOOL GetCachedRsaKey(TPMT_PUBLIC*    publicArea,
		     TPMT_SENSITIVE* sensitive,
		     RAND_STATE*     rand  // IN: if not NULL, the deterministic
		     //     RNG state
		     )
{
    int keyBits = publicArea->parameters.rsaDetail.keyBits;
    int index;
    //
    if(KeyCacheLoaded(publicArea, sensitive, rand))
	{
	    for(index = 0; index < RSA_KEY_CACHE_ENTRIES; index++)
		{
		    if((s_rsaKeyCache[index].publicModulus.t.size * 8) == keyBits)
			{
			    publicArea->unique.rsa   = s_rsaKeyCache[index].publicModulus;
			    sensitive->sensitive.rsa = s_rsaKeyCache[index].privateExponent;
			    return TRUE;
			}
		}
	    return FALSE;
	}
    return s_keyCacheLoaded;
}
#endif  // defined SIMULATION && defined USE_RSA_KEY_CACHE
